package com.seanLab.tool.TagSuggestModel.ExpandRank.common;


import com.seanLab.dto.SuggestModelArticleDto;
import com.seanLab.tool.TagSuggestModel.ExpandRank.common.TagFilter;
import com.seanLab.tool.TagSuggestModel.ExpandRank.common.WordFeatureExtractor;
//import org.apache.lucene.analysis.WhitespaceAnalyzer;
import org.apache.lucene.analysis.core.WhitespaceAnalyzer;
import org.apache.lucene.document.Document;
import org.apache.lucene.document.Field;
import org.apache.lucene.document.FieldType;
import org.apache.lucene.document.TextField;
import org.apache.lucene.index.IndexOptions;
import org.apache.lucene.index.IndexWriter;
import org.apache.lucene.index.IndexWriterConfig;
import org.apache.lucene.store.Directory;
import org.apache.lucene.store.FSDirectory;
import org.thunlp.misc.StringUtil;
import org.thunlp.text.Lexicon;

import java.io.File;
import java.io.IOException;
import java.nio.file.Paths;
import java.util.HashSet;
import java.util.List;
import java.util.Properties;
import java.util.Set;
import java.util.logging.Logger;

/**
 * 训练ExpandRank
 * 建立文章Lucene索引，以快速找到最相似的若干篇文章
 */
public class TrainExpandRank {
    private static Logger LOG = Logger.getAnonymousLogger();
    private WordFeatureExtractor extractor;
    private TagFilter tagFilter = null;
    private Properties config = null;
    private Lexicon wordlex = null;
    private Lexicon taglex = null;
    private String fold = "1";
    private int minTagFreq = 1;
    private Set<String> filtered = new HashSet<String>();


//    public void run(String[] args) throws Exception {
//        Flags flags = new Flags();
//        flags.add("input");
//        flags.add("output");
//        flags.add("config");
//        flags.parseAndCheck(args);
//
//        Properties config = ConfigIO.configFromString(flags.getString("config"));
//
//        train(flags.getString("input"), flags.getString("output"), config);
//    }

    public String train(List<SuggestModelArticleDto> articles, String workingDir, Properties config) throws IOException {
        this.config = config;
        minTagFreq = Integer.parseInt(config.getProperty("mintagfreq", "1"));
        wordlex = new Lexicon();
        taglex = new Lexicon();
        WordFeatureExtractor.buildLexicons(articles, workingDir, wordlex, taglex, config);
        WordFeatureExtractor e  = new WordFeatureExtractor(config);
        e.setWordLexicon(wordlex);
        e.setTagLexicon(taglex);
        extractor = e;
        tagFilter = new TagFilter(config, taglex);
        String modelPath = workingDir + File.separator + "model_" + articles.size();
        buildIndexes(articles, new File(modelPath));
        return modelPath;
    }

//    public void train(String input, String modelPath, Properties config)
//            throws IOException {
//        this.config = config;
//        this.fold = config.getProperty("fold", "");
//        minTagFreq = Integer.parseInt(config.getProperty("mintagfreq", "1"));
//        wordlex = new Lexicon();
//        taglex = new Lexicon();
//        WordFeatureExtractor.buildLexicons(
//                input, wordlex, taglex, config);
//        WordFeatureExtractor e = new WordFeatureExtractor(config);
//        e.setTagLexicon(taglex);
//        e.setWordLexicon(wordlex);
//        extractor = e;
//        tagFilter = new TagFilter(config, taglex);
//        buildIndexes(input, new File(modelPath));
//    }

    public void buildIndexes(List<SuggestModelArticleDto> articles, File modelDir) throws IOException {
        if (!modelDir.exists()) {
            modelDir.mkdir();
        }
//        //for lucene 2.0
//        WhitespaceAnalyzer analyzer = new WhitespaceAnalyzer();
//        IndexWriter docsIndex = new IndexWriter(new File(modelDir, "docs"), analyzer);
//        //for lucene 2.0

        //for lucene 6.6
        Directory dir = FSDirectory.open(Paths.get(modelDir.getPath(), "docs"));
        IndexWriter docsIndex = new IndexWriter(dir, new IndexWriterConfig(new WhitespaceAnalyzer()));
        //for lucene 6.6
        for (SuggestModelArticleDto article : articles) {
            Document contentDoc = makeContentDoc(article);
            docsIndex.addDocument(contentDoc);
        }
        LOG.info("Added " + articles.size() + " documents.");
//        LOG.info("Optimizing docs index..."); //for lucene 2.0
//        docsIndex.optimize(); //for lucene2.0
        docsIndex.close();
    }

//    public void buildIndexes(String input, File modelDir)
//            throws IOException {
//        if (!modelDir.exists()) {
//            modelDir.mkdir();
//        }
//
//        Set<String> whitelist = new HashSet<String>();
//        Set<String> blacklist = new HashSet<String>();
//
//        if (config.getProperty("whitelist", "").length() > 0) {
//            whitelist.addAll(
//                    Arrays.asList(config.getProperty("whitelist", "").split(",")));
//        }
//        if (config.getProperty("blacklist", "").length() > 0) {
//            blacklist.addAll(
//                    Arrays.asList(config.getProperty("blacklist", "").split(",")));
//        }
//
//        WhitespaceAnalyzer analyzer = new WhitespaceAnalyzer();
//        JsonUtil J = new JsonUtil();
//        IndexWriter docsIndex =
//                new IndexWriter(new File(modelDir, "docs"), analyzer);
//
//        RecordReader reader = new RecordReader(input);
//        while (reader.next()) {
//            //Post p = J.fromJson(reader.value(), Post.class);
//            KeywordPost p = J.fromJson(reader.value(), KeywordPost.class);
////            if (blacklist.contains(p.getUserId())) {
////                continue;
////            }
////            if (whitelist.size() > 0 && !whitelist.contains(p.getUserId())) {
////                continue;
////            }
//            if (fold.length() > 0 && p.getExtras().equals(fold)) {
//                continue;
//            }
//            Document contentDoc = makeContentDoc(p);
//            docsIndex.addDocument(contentDoc);
//            if (reader.numRead() % 5000 == 0) {
//                LOG.info("Added " + reader.numRead() + " documents.");
//            }
//        }
//        reader.close();
//
//        LOG.info("Optimizing docs index...");
//        docsIndex.optimize();
//        docsIndex.close();
//    }

    public Document makeContentDoc(SuggestModelArticleDto article) {
        String[] words = extractor.extractKeyword(article, true, true);
        String docString = StringUtil.join(words, " ");
        Document d = new Document();
//        d.add(new Field("articleID", article.getArticleID(),
//                Field.Store.YES, Field.Index.UN_TOKENIZED));

//        //for lucene2.0
//        d.add(new Field("content", docString,
//                Field.Store.YES, Field.Index.TOKENIZED));
//        tagFilter.filter(article.getTags(), filtered);
//        d.add(new Field("tags", StringUtil.join(filtered, " "),
//                Field.Store.YES, Field.Index.TOKENIZED));
//        //for lucene2.0

        //for lucene 6.6
        FieldType fieldType = new FieldType();
//        fieldType.setStoreTermVectors(true);
        fieldType.setIndexOptions(IndexOptions.DOCS_AND_FREQS);
        fieldType.setTokenized(true);
        fieldType.setStored(true);
        fieldType.freeze();
        d.add(new Field("content", docString, fieldType)); //for lucene6.6
        tagFilter.filter(article.getTags(), filtered);
        d.add(new Field("tags", StringUtil.join(filtered, " "), fieldType)); //for lucene6.6

//        d.add(new TextField("content", docString, Field.Store.YES)); //for lucene6.6
//        tagFilter.filter(article.getTags(), filtered);
//        d.add(new TextField("tags", StringUtil.join(filtered, " "), Field.Store.YES)); //for lucene6.6
        //for lucene 6.6
        return d;
    }

//    public Document makeContentDoc(KeywordPost p) {
//        //String [] words = extractor.extract(p);
//        String[] words = extractor.extractKeyword(p, true, true, true);
//        String docString = StringUtil.join(words, " ");
//        Document d = new Document();
//        d.add(new Field("doc_id", p.getArticleID(),
//                Field.Store.YES, Field.Index.UN_TOKENIZED));
//        d.add(new Field("content", docString,
//                Field.Store.YES, Field.Index.TOKENIZED));
//        tagFilter.filter(p.getTags(), filtered);
//        d.add(new Field("tags", StringUtil.join(filtered, " "),
//                Field.Store.YES, Field.Index.TOKENIZED));
////        d.add(new Field("user_id", p.getUserId(),
////                Field.Store.YES, Field.Index.UN_TOKENIZED));
//        return d;
//    }

}
